# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from tqdm import tqdm
import xml.etree.ElementTree as ET
from sacremoses import MosesTokenizer
from pathlib import Path
import argparse


def moses_cut(in_file, out_file, lang):  # in_file: 输入文件路径,out_file: 输出文件路径,lang: 语言
    # 创建一个MosesTokenizer对象，用于对文本进行分词处理，lang参数指定了分词时所用的语言（如'en'表示英文，'de'表示德文）
    # MosesTokenizer类是用于对自然语言文本进行分词的工具，支持多种语言
    # 该类能够将连续的文本字符串按照语言学规则切分成单词或标记（token）
    # 这里通过指定lang参数，MosesTokenizer会自动适配不同语言的分词规范
    # 例如英文会处理缩写、标点，德文会处理复合词等
    # 该分词器广泛应用于机器翻译、文本预处理等NLP任务
    mt = MosesTokenizer(lang=lang)
    
    out_f = open(out_file, "w", encoding="utf8") # 创建输出文件，w模式为写入，编码为utf8,如果文件不存在，则创建

    with open(in_file, "r", encoding="utf8") as f:

        for line in f.readlines():  # readlines()方法读取文件内容并返回一个列表，列表中的元素是文件中的每一行
            line = line.strip()  # 去除行首和行尾的空白字符
            if not line: # 判断行是否为空
                continue # 跳过空行
            cut_line = mt.tokenize(line, return_str=True)  # 分词,return_str=True表示返回字符串
            out_f.write(cut_line.lower() + "\n")  # 变为小写，并写入文件
    out_f.close()


if __name__ == "__main__":
    # 下面这行代码创建了一个命令行参数解析器对象parser
    # argparse.ArgumentParser()是Python标准库argparse中的一个类
    # 这个解析器的作用是帮助我们从命令行读取和解析用户输入的参数和选项
    # 通过parser对象，我们可以定义程序支持哪些命令行参数（比如输入文件路径、输出目录、语言类型等）
    # 这样用户在运行脚本时可以灵活地指定不同的参数，而不需要修改代码
    # 解析器还会自动生成帮助信息，并在参数输入不合法时给出提示
    parser = argparse.ArgumentParser()  # 创建参数解析器对象
    # 添加参数,-p,--pair_dir 参数，表示输入文件路径
    parser.add_argument(
        "-p",
        "--pair_dir",
        default=None,
        type=str,
        help="The directory which contains language pair files.",
    )
    # 添加参数，-d,--dest_dir参数，表示保存处理后的文件路径
    parser.add_argument(
        "-d",
        "--dest_dir",
        default=None,
        type=str,
        help="The destination directory to save processed train, dev and test file.",
    )

    # 添加参数
    parser.add_argument("--src_lang", default="de", type=str, help="source language")
    parser.add_argument("--trg_lang", default="en", type=str, help="target language")
    # 解析参数
    args = parser.parse_args()  # 解析参数，args是一个列表，包含了传递的参数值
    print(args)
    if not args.pair_dir:  # 如果不传参，就抛异常
        raise ValueError("Please specify --pair_dir")

    if not os.path.exists(args.dest_dir):  # 检查目标目录是否存在
        os.makedirs(args.dest_dir)         # 如果不存在则创建目标目录
        
    local_data_path = Path(args.pair_dir)  # 获取本地数据路径
    data_dir = Path(args.dest_dir)  # 获取保存路径

    # 分词
    # 遍历数据集的三种模式：训练集、验证集和测试集
    for mode in ["train", "val", "test"]:
        # 对每种模式下的源语言文件进行分词处理
        # 调用moses_cut函数对源语言文件进行分词处理
        # local_data_path / f"{mode}.{args.src_lang}" 构建当前模式（如train、val、test）下源语言文件的完整路径
        # data_dir / f"{mode}_src.cut.txt" 构建分词后结果要保存的目标文件路径
        # lang=args.src_lang 指定分词器的语言为源语言
        # 使用参数
        moses_cut(
            local_data_path / f"{mode}.{args.src_lang}",  # 获取源语言原始文本文件路径
            data_dir / f"{mode}_src.cut.txt",             # 指定分词后文本的保存路径
            lang=args.src_lang,                           # 设置分词器的语言
        )
        # 分词完成后，输出提示信息，表明当前模式下的源语言文本已经分词完毕
        print(f"[{mode}] 源语言文本分词完成")


        moses_cut(
            local_data_path / f"{mode}.{args.trg_lang}",  # 读取目标语言文件
            data_dir / f"{mode}_trg.cut.txt",
            lang=args.trg_lang,
        )
        print(f"[{mode}] 目标语言文本分词完成")
        
    # 创建文件夹，移动读取的文本到刚创建的文件夹里
    # if not data_dir.exists():
    #     data_dir.mkdir(parents=True)
    # for fpath in local_data_path.glob("*.txt"): # 遍历所有分词后的文件,并移动到目标文件夹
    #     fpath.rename(data_dir / fpath.name)
